{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a861fa18",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn, optim\n",
    "from torchvision import models, transforms, datasets\n",
    "from PIL import Image\n",
    "import mltools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3be045c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tfrs(transform, path=\"./data\", batch_size=8, num_workers=8, pin_memory=True, drop_last=True):\n",
    "    data = datasets.ImageFolder(f\"{path}/TFRS\", transform=transform)\n",
    "    train_data, val_data, test_data = mltools.split_data(data, [0.7, 0.15, 0.15])\n",
    "    return mltools.iter_data([train_data, val_data, test_data], batch_size, num_workers=num_workers, pin_memory=pin_memory, drop_last=drop_last)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d853916",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\")\n",
    "transform = transforms.Compose([transforms.Resize((256, 256)), transforms.ToTensor(), transforms.ConvertImageDtype(torch.float32), transforms.Normalize(mean=0.5, std=0.5)])\n",
    "train_iter, val_iter, test_iter = tfrs(transform)  # 获取训练集、验证集、测试集\n",
    "labels = train_iter.dataset.dataset.classes\n",
    "model = models.resnet18(num_classes=len(labels))  # 设置模型结构\n",
    "model.to(device)\n",
    "loss = nn.CrossEntropyLoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr=1e-2)\n",
    "ml = mltools.MachineLearning(\"TFRS\")\n",
    "ml.add_model(model)\n",
    "epoch, timer, recorder = ml.batch_create()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79f4df45",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 训练\n",
    "num_epochs = epoch(3)\n",
    "animator = ml.create_animator(xlabel=\"epoch\", xlim=[0, epoch.totol_epoch + 1], ylim=-0.1, legend=[\"train loss\", \"train acc\", \"val acc\"])  # 创建动画器\n",
    "for current_epoch in range(1, num_epochs + 1):\n",
    "    timer.start()  # 开始计时\n",
    "\n",
    "    # 计算训练集\n",
    "    metric_train = mltools.Accumulator(3)  # 累加器：(train_loss, train_acc, train_size)\n",
    "    model.train()  # 训练模式\n",
    "    for x, y in train_iter:\n",
    "        x, y = x.to(device), y.to(device)  # 转换x、y\n",
    "        y_train = model(x)  # 计算模型\n",
    "        train_loss = loss(y_train, y)  # 计算损失\n",
    "        train_pred = y_train.argmax(dim=1)  # 计算预测值\n",
    "        train_acc = (train_pred == y).sum()  # 计算准确数\n",
    "\n",
    "        # 梯度更新\n",
    "        optimizer.zero_grad()\n",
    "        train_loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        metric_train.add(train_loss * y.numel(), train_acc, y.numel())  # 累加损失、准确数、样本数\n",
    "    recorder[0].append(metric_train[0] / metric_train[2])  # 计算平均损失\n",
    "    recorder[1].append(metric_train[1] / metric_train[2])  # 计算准确率\n",
    "\n",
    "    # 计算验证集\n",
    "    metric_val = mltools.Accumulator(2)  # 累加器：(val_acc, val_size)\n",
    "    model.eval()  # 验证模式\n",
    "    with torch.no_grad():\n",
    "        for x, y in val_iter:\n",
    "            x, y = x.to(device), y.to(device)  # 转换x、y\n",
    "            y_val = model(x)  # 计算模型\n",
    "            val_acc = (y_val.argmax(dim=1) == y).sum()  # 计算准确数\n",
    "            metric_val.add(val_acc, y.numel())  # 累加准确数、样本数\n",
    "    recorder[2].append(metric_val[0] / metric_val[1])  # 计算准确率\n",
    "\n",
    "    timer.stop()  # 停止计时\n",
    "\n",
    "    # 打印输出值\n",
    "    ml.logger.info(f\"train loss {recorder[0][-1]:.3f}, train acc {recorder[1][-1]:.3f}, val acc {recorder[2][-1]:.3f}\")\n",
    "    ml.print_training_time_massage(timer, num_epochs, current_epoch)\n",
    "    ml.logger.info(f\"trained on {str(device)}\")\n",
    "    animator.show(recorder.data)\n",
    "else:\n",
    "    # 打印输出值\n",
    "    ml.logger.info(f\"train loss {recorder[0][-1]:.3f}, train acc {recorder[1][-1]:.3f}, val acc {recorder[2][-1]:.3f}\")\n",
    "    ml.print_training_time_massage(timer, num_epochs, current_epoch)\n",
    "    ml.logger.info(f\"trained on {str(device)}\")\n",
    "    animator.show(recorder.data)\n",
    "ml.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e34baf75",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 测试\n",
    "model.eval()\n",
    "metric = mltools.Accumulator(2)  # 累加器：(test_acc, test_size)\n",
    "with torch.no_grad():\n",
    "    for x, y in test_iter:\n",
    "        x, y = x.to(device), y.to(device)  # 转换x、y\n",
    "        y_test = model(x)  # 计算模型\n",
    "        test_pred = y_test.argmax(dim=1)  # 计算预测值\n",
    "        test_acc = (test_pred == y).sum()  # 计算准确数\n",
    "        metric.add(test_acc, y.numel())  # 累加准确数、样本数\n",
    "ml.logger.info(f\"test acc {metric[0] / metric[1]:.3f}\")  # 计算准确率并输出"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95ff4db6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 预测\n",
    "image = \"../data/TFRS/仿宋/仿宋_01.png\"\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    image_tensor = transform(Image.open(image).convert(\"RGB\")).unsqueeze(0)\n",
    "    image_tensor = image_tensor.to(device)\n",
    "    print(labels[model(image_tensor).argmax(dim=1)[0]])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tfrs",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
